
package com.simple.mqtt.server.session;

import com.simple.mqtt.common.domain.TopicFilter;
import com.simple.mqtt.common.task.MqttPendingPublish;
import com.simple.mqtt.common.task.MqttPendingQos2Publish;
import com.simple.mqtt.server.domain.Subscribe;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * In-memory MQTT session manager implementation
 * Uses thread-safe collections like ConcurrentHashMap for session information storage
 * Supports all IMqttSessionManager interface functions
 *
 * @author L.cm
 */
public class InMemoryMqttSessionManager implements IMqttSessionManager {
    /**
     * Message ID storage clientId: messageId
     */
    private final ConcurrentMap<String, AtomicInteger> messageIdStore = new ConcurrentHashMap<>();
    /**
     * Subscription storage, supports shared subscriptions
     */
    private final TrieTopicManager topicManager = new TrieTopicManager();
    /**
     * QoS1 message process storage clientId: {msgId: Object}
     */
    private final ConcurrentMap<String, ConcurrentMap<Integer, MqttPendingPublish>> pendingPublishStore = new ConcurrentHashMap<>();
    /**
     * QoS2 message process storage clientId: {msgId: Object}
     */
    private final ConcurrentMap<String, ConcurrentMap<Integer, MqttPendingQos2Publish>> pendingQos2PublishStore = new ConcurrentHashMap<>();

    @Override
    public void addSubscribe(TopicFilter topicFilter, String clientId, int mqttQoS) {
        topicManager.addSubscribe(topicFilter, clientId, (short) mqttQoS);
    }

    @Override
    public void removeSubscribe(String topicFilter, String clientId) {
        topicManager.removeSubscribe(topicFilter, clientId);
    }

    /**
     * 移除客户端的所有订阅关系
     * @param clientId 客户端ID
     */
    public void removeSubscribe(String clientId) {
        topicManager.removeSubscribe(clientId);
    }

    @Override
    public Byte searchSubscribe(String topicName, String clientId) {
        return topicManager.searchSubscribe(topicName, clientId);
    }

    @Override
    public List<Subscribe> searchSubscribe(String topicName) {
        return topicManager.searchSubscribe(topicName);
    }

    @Override
    public List<Subscribe> getSubscriptions(String clientId) {
        return topicManager.getSubscriptions(clientId);
    }

    @Override
    public void addPendingPublish(String clientId, int messageId, MqttPendingPublish pendingPublish) {
        Map<Integer, MqttPendingPublish> data = pendingPublishStore.computeIfAbsent(clientId, (key) -> new ConcurrentHashMap<>(16));
        data.put(messageId, pendingPublish);
    }

    @Override
    public MqttPendingPublish getPendingPublish(String clientId, int messageId) {
        Map<Integer, MqttPendingPublish> data = pendingPublishStore.get(clientId);
        if (data == null) {
            return null;
        }
        return data.get(messageId);
    }

    @Override
    public void removePendingPublish(String clientId, int messageId) {
        Map<Integer, MqttPendingPublish> data = pendingPublishStore.get(clientId);
        if (data != null) {
            data.remove(messageId);
        }
    }

    @Override
    public void addPendingQos2Publish(String clientId, int messageId, MqttPendingQos2Publish pendingQos2Publish) {
        Map<Integer, MqttPendingQos2Publish> data = pendingQos2PublishStore.computeIfAbsent(clientId, (key) -> new ConcurrentHashMap<>(16));
        data.put(messageId, pendingQos2Publish);
    }

    @Override
    public MqttPendingQos2Publish getPendingQos2Publish(String clientId, int messageId) {
        Map<Integer, MqttPendingQos2Publish> data = pendingQos2PublishStore.get(clientId);
        if (data == null) {
            return null;
        }
        return data.get(messageId);
    }

    @Override
    public void removePendingQos2Publish(String clientId, int messageId) {
        Map<Integer, MqttPendingQos2Publish> data = pendingQos2PublishStore.get(clientId);
        if (data != null) {
            data.remove(messageId);
        }
    }

    @Override
    public int getPacketId(String clientId) {
        // 为每个客户端创建独立的消息ID生成器，避免冲突
        // 当ID达到最大值(0xffff)时重置为1
        AtomicInteger packetIdGen = messageIdStore.computeIfAbsent(clientId, (key) -> new AtomicInteger(1));
        return packetIdGen.getAndUpdate(current -> (current % 0xffff) == 0 ? 1 : current + 1);
    }

    @Override
    public boolean hasSession(String clientId) {
        return pendingQos2PublishStore.containsKey(clientId)
                || pendingPublishStore.containsKey(clientId)
                || messageIdStore.containsKey(clientId)
                || !topicManager.getSubscriptions(clientId).isEmpty();
    }

    @Override
    public boolean expire(String clientId, int sessionExpirySeconds) {
        return false;
    }

    @Override
    public boolean active(String clientId) {
        return false;
    }

    @Override
    public void remove(String clientId) {
        removeSubscribe(clientId);
        pendingPublishStore.remove(clientId);
        pendingQos2PublishStore.remove(clientId);
        messageIdStore.remove(clientId);
    }

    @Override
    public void clean() {
        topicManager.clear();
        pendingPublishStore.clear();
        pendingQos2PublishStore.clear();
        messageIdStore.clear();
    }

}
